import base64
import os
import time

import requests

# 服务地址
url = "http://10.157.223.52:9002/ocr"  # 使用正确的接口路径

# 输入文件路径，可为图片或 PDF
file_path = r"C:\Users\L\PycharmProjects\PaddleOCR\合同_JS66642025000323.pdf"


def call_ocr(file_bytes: bytes, file_type: int, visualize=True):
    """将文件字节发送至 OCR 服务并返回响应。"""
    base64_data = base64.b64encode(file_bytes).decode("utf-8")
    payload = {
        "file": base64_data,
        "fileType": file_type,  # 0: PDF, 1: 图像
        "visualize": visualize,
    }
    response = requests.post(url, json=payload)
    response.raise_for_status()
    return response.json()


def process_pdf(pdf_path: str):
    """直接读取 PDF 字节并以 fileType=0 方式发送。"""
    with open(pdf_path, "rb") as f:
        pdf_bytes = f.read()
    print("PDF 文件读取完成，开始发送 OCR 请求。")
    return call_ocr(pdf_bytes, file_type=0, visualize=True)


def process_image(image_path: str):
    with open(image_path, "rb") as f:
        image_bytes = f.read()
    return call_ocr(image_bytes, file_type=1, visualize=True)


def _extract_ocr_results(payload: dict):
    """递归寻找 payload 中的 ocrResults 列表。"""
    if not isinstance(payload, dict):
        return []
    if isinstance(payload.get("ocrResults"), list):
        return payload["ocrResults"]
    inner = payload.get("result")
    if isinstance(inner, dict):
        return _extract_ocr_results(inner)
    return []


def collect_rec_texts(result_data):
    """从 OCR 响应中汇总 rec_texts。"""
    texts = []
    entries = result_data if isinstance(result_data, list) else [result_data]
    for entry in entries:
        if not isinstance(entry, dict):
            continue
        ocr_results = _extract_ocr_results(entry.get("result", entry))
        for ocr_item in ocr_results:
            pruned = ocr_item.get("prunedResult", ocr_item)
            texts.extend(pruned.get("rec_texts", []))
    return texts


def save_ocr_images(result_data):
    """保存OCR结果中的图片。"""
    saved_count = 0
    entries = result_data if isinstance(result_data, list) else [result_data]
    for entry in entries:
        if not isinstance(entry, dict):
            continue

        # 检查直接在结果中的ocrResults
        if isinstance(entry.get("ocrResults"), list):
            for i, res in enumerate(entry["ocrResults"]):
                if "ocrImage" in res:
                    ocr_img_path = f"ocr_{saved_count}.jpg"
                    with open(ocr_img_path, "wb") as f:
                        f.write(base64.b64decode(res["ocrImage"]))
                    print(f"Output image saved at {ocr_img_path}")
                    saved_count += 1

        # 检查在result字段中的ocrResults
        if isinstance(entry.get("result"), dict):
            inner_result = entry["result"]
            if isinstance(inner_result.get("ocrResults"), list):
                for i, res in enumerate(inner_result["ocrResults"]):
                    if "ocrImage" in res:
                        ocr_img_path = f"ocr_{saved_count}.jpg"
                        with open(ocr_img_path, "wb") as f:
                            f.write(base64.b64decode(res["ocrImage"]))
                        print(f"Output image saved at {ocr_img_path}")
                        saved_count += 1
    return saved_count


def main():
    start_time = time.time()
    _, ext = os.path.splitext(file_path)
    ext = ext.lower()

    try:
        if ext == ".pdf":
            result = process_pdf(file_path)
        else:
            result = process_image(file_path)
    except requests.HTTPError as exc:
        print("请求失败:", exc.response.status_code)
        print("错误信息:", exc.response.text)
        return

    rec_texts = collect_rec_texts(result)

    # 保存OCR结果图片
    saved_images = save_ocr_images(result)
    if saved_images > 0:
        print(f"成功保存了 {saved_images} 张OCR结果图片")

    print("运行时间：", time.time() - start_time)
    if rec_texts:
        print("识别文本:", rec_texts)
    # print("OCR 结果:", result)


if __name__ == "__main__":
    main()
